'''转换原始标签格式至cholect50-crossval格式，将原先的id标签转换为onehot的
'''
import pathlib
import glob
import os
import json
import numpy as np
from tqdm import tqdm

# 控制在合并同一帧的多个三元组标签时是否多次累加同一component的标签
# 若选择多次叠加，得到的0、1标签中可能会有超过1的
OVERFLOW_ONE = False


data_root = pathlib.Path("/opt/data/private/dataset/CholecT50")
origin_lb_path = data_root / "labels"
triplet_lb_path = data_root / "triplet"
instrument_lb_path = data_root / "instrument"
verb_lb_path = data_root / "verb"
target_lb_path = data_root / "target"
phase_lb_path = data_root / "phase"

os.makedirs(triplet_lb_path, exist_ok=True)
os.makedirs(instrument_lb_path, exist_ok=True)
os.makedirs(verb_lb_path, exist_ok=True)
os.makedirs(target_lb_path, exist_ok=True)
os.makedirs(phase_lb_path, exist_ok=True)

origin_json_file_list = glob.glob(f"{origin_lb_path}/*.json")
# 逐个视频标注处理
for json_file_path in tqdm(origin_json_file_list):
    triplet_lbs = []
    instrument_lbs = []
    verb_lbs = []
    target_lbs = []
    phase_lbs = []

    # 从CholecT50的json标注文件中提取各component的标注
    with open(json_file_path, "r", encoding="utf-8") as origin_label_file:
        json_data = json.load(origin_label_file)
        video_id: int = json_data['video']
        annotations = json_data['annotations']
        num_frames = int(json_data['num_frames'])
        for frame_idx in range(num_frames):
            # 一帧的标注列表
            annotations_frame: list = annotations[str(frame_idx)]
            instrument_lb = np.zeros([6])
            verb_lb = np.zeros([10])
            target_lb = np.zeros([15])
            triplet_lb = np.zeros([100])
            phase_lb = np.zeros([7])
            # 叠加某一帧的所有标注
            for anno in annotations_frame:
                if anno[0] >= 0:
                    triplet_lb[anno[0]] += 1
                if anno[1] >= 0:
                    instrument_lb[anno[1]] += 1
                if anno[7] >= 0:
                    verb_lb[anno[7]] += 1
                if anno[8] >= 0:
                    target_lb[anno[8]] += 1
                if anno[14] >= 0:
                    phase_lb[anno[14]] += 1
            # 若需要不超过1，则将所有有标注的类别都转为1
            if not OVERFLOW_ONE:
                triplet_lb[triplet_lb > 0] = 1
                instrument_lb[instrument_lb > 0] = 1
                verb_lb[verb_lb > 0] = 1
                target_lb[target_lb > 0] = 1
                phase_lb[phase_lb > 0] = 1

            # 在每条记录的前面添加帧号，方便后续直接读取帧图片
            triplet_lbs.append(np.concatenate([[frame_idx], triplet_lb]))
            instrument_lbs.append(np.concatenate([[frame_idx], instrument_lb]))
            verb_lbs.append(np.concatenate([[frame_idx], verb_lb]))
            target_lbs.append(np.concatenate([[frame_idx], target_lb]))
            phase_lbs.append(np.concatenate([[frame_idx], phase_lb]))
    
    np.savetxt(
        triplet_lb_path / f"VID{video_id:02d}.txt",
        np.array(triplet_lbs, dtype=np.int32),
        fmt="%i",
        delimiter=","
    )
    np.savetxt(
        instrument_lb_path / f"VID{video_id:02d}.txt",
        np.array(instrument_lbs, dtype=np.int32),
        fmt="%i",
        delimiter=","
    )
    np.savetxt(
        verb_lb_path / f"VID{video_id:02d}.txt",
        np.array(verb_lbs, dtype=np.int32),
        fmt="%i",
        delimiter=","
    )
    np.savetxt(
        target_lb_path / f"VID{video_id:02d}.txt",
        np.array(target_lbs, dtype=np.int32),
        fmt="%i",
        delimiter=","
    )
    np.savetxt(
        phase_lb_path / f"VID{video_id:02d}.txt",
        np.array(phase_lbs, dtype=np.int32),
        fmt="%i",
        delimiter=","
    )
